"""
模糊检测器实现类

基于分割模型实现模糊检测功能，支持单张图片和批量检测。
"""

import os
import sys
import numpy as np
from PIL import Image
from typing import Dict, Any, Optional, Union, List
from pathlib import Path

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference


class BlurDetectorImpl:
    """
    模糊检测器实现类
    
    基于分割模型实现模糊检测功能，可以检测图像中的模糊区域。
    返回1通道的二值图，其中1表示模糊区域，0表示非模糊区域。
    """
    
    def __init__(self, 
                 model_path: str, 
                 device: str = 'cpu',
                 blur_class_id: int = 1,
                 confidence_threshold: float = 0.5):
        """
        初始化模糊检测器
        
        Args:
            model_path (str): 分割模型文件路径
            device (str): 推理设备，'cpu' 或 'cuda'
            blur_class_id (int): 模糊类别ID，默认1
            confidence_threshold (float): 置信度阈值，默认0.5
        """
        self.model_path = model_path
        self.device = device
        self.blur_class_id = blur_class_id
        self.confidence_threshold = confidence_threshold
        
        # 创建推理器实例
        try:
            self.infer = create_inference(model_path, device=device)
            print(f"模糊检测器初始化成功，使用设备: {device}")
        except Exception as e:
            raise RuntimeError(f"初始化模糊检测器失败: {str(e)}")
    
    def detect_blur_single(self, image: Union[str, Image.Image, np.ndarray], 
                          output_path: Optional[str] = None) -> Dict[str, Any]:
        """
        检测单张图片中的模糊区域
        
        Args:
            image: 图像数据，可以是图像路径(str)、PIL Image对象或numpy数组
            output_path (str, optional): 输出二值图路径，如果提供则保存二值图
            
        Returns:
            Dict[str, Any]: 检测结果，包含模糊区域二值图
        """
        try:
            # 记录原始图像尺寸
            original_size = None
            
            # 处理不同类型的输入
            if isinstance(image, str):
                # 如果是路径，加载图像并记录尺寸
                pil_image = Image.open(image)
                original_size = pil_image.size
                result = self.infer.infer_single(image)
            elif isinstance(image, Image.Image):
                # 如果是PIL Image，记录尺寸
                original_size = image.size
                result = self.infer.infer_single(image)
            elif isinstance(image, np.ndarray):
                # 如果是numpy数组，直接传递给推理器（base_inference已处理格式转换）
                # 记录原始图像尺寸
                if len(image.shape) == 2:
                    original_size = (image.shape[1], image.shape[0])  # (width, height)
                elif len(image.shape) == 3 and image.shape[2] == 3:
                    original_size = (image.shape[1], image.shape[0])  # (width, height)
                else:
                    raise ValueError(f"不支持的numpy数组形状: {image.shape}，期望形状为(H, W)或(H, W, 3)")
                
                result = self.infer.infer_single(image)
            else:
                raise ValueError(f"不支持的图像类型: {type(image)}")
            
            # 从分割结果中提取模糊区域
            blur_mask = self._extract_blur_mask(result['segmentation_mask'])
            
            # 如果原始图像尺寸已知，将二值图resize回原图尺寸
            if original_size is not None:
                blur_mask_resized = self._resize_mask_to_original(blur_mask, original_size)
            else:
                blur_mask_resized = blur_mask
            
            # 构建返回结果
            detection_result = {
                'blur_mask': blur_mask_resized,  # 1通道二值图，1表示模糊区域
                'original_size': original_size,
                'mask_size': blur_mask_resized.shape,
                'num_blur_pixels': int(np.sum(blur_mask_resized)),
                'blur_ratio': float(np.sum(blur_mask_resized)) / (blur_mask_resized.shape[0] * blur_mask_resized.shape[1]),
                'segmentation_mask': result['segmentation_mask'],
                'num_classes': result['num_classes'],
                'image_shape': result['segmentation_mask'].shape,
                # 添加完整的模型推理结果
                'model_result': result
            }
            
            # 如果提供了输出路径，保存二值图
            if output_path:
                self._save_blur_mask(blur_mask_resized, output_path)
            
            return detection_result
            
        except Exception as e:
            raise RuntimeError(f"模糊检测失败: {str(e)}")
    
    
    def _extract_blur_mask(self, segmentation_mask: np.ndarray) -> np.ndarray:
        """
        从分割掩码中提取模糊区域，生成1通道二值图
        
        Args:
            segmentation_mask (np.ndarray): 分割掩码
            
        Returns:
            np.ndarray: 1通道二值图，1表示模糊区域，0表示非模糊区域
        """
        # 找到标签为模糊类别的像素
        blur_pixels = (segmentation_mask == self.blur_class_id)
        
        # 转换为uint8格式的二值图
        blur_mask = (blur_pixels * 255).astype(np.uint8)
        
        return blur_mask
    
    def _resize_mask_to_original(self, mask: np.ndarray, original_size: tuple) -> np.ndarray:
        """
        将二值图resize回原图尺寸
        
        Args:
            mask (np.ndarray): 二值图
            original_size (tuple): 原始图像尺寸 (width, height)
            
        Returns:
            np.ndarray: resize后的二值图
        """
        import cv2
        
        # 使用最近邻插值保持二值特性
        resized_mask = cv2.resize(
            mask, 
            original_size, 
            interpolation=cv2.INTER_NEAREST
        )
        
        return resized_mask
    
    def _save_blur_mask(self, blur_mask: np.ndarray, output_path: str):
        """
        保存模糊检测二值图
        
        Args:
            blur_mask (np.ndarray): 二值图
            output_path (str): 输出路径
        """
        try:
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 保存为PNG格式的二值图
            pil_mask = Image.fromarray(blur_mask, mode='L')
            pil_mask.save(output_path)
            
            print(f"模糊检测二值图已保存到: {output_path}")
            
        except Exception as e:
            print(f"保存二值图失败: {str(e)}")
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        try:
            return self.infer.get_model_info()
        except Exception as e:
            return {'error': f"获取模型信息失败: {str(e)}"}
    
    def validate_model(self) -> bool:
        """
        验证模型是否有效
        
        Returns:
            bool: 模型是否有效
        """
        try:
            return self.infer.validate_model()
        except Exception as e:
            print(f"模型验证失败: {str(e)}")
            return False
    
    def get_task_type(self) -> str:
        """
        获取任务类型
        
        Returns:
            str: 任务类型
        """
        try:
            return self.infer.get_task_type()
        except Exception as e:
            return f"获取任务类型失败: {str(e)}"
